from copy import deepcopy
import random
import time
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
from dag import Circuit_DAG, CNOT_Node

from tool import draw_a_dag


class QPU:
    # 构建一个量子芯片chiplet
    # chip_type - Chip指向的类型，有 "square", "hexagon", "heavy_square", "heavy_hexagon" 四种
    # inner_chip_size - 芯片粒的大小，有 6, 7, 8, 9 四种
    # inter_chip_array - 芯片组数，有 [2, 2], [2, 3], [3, 3], [3, 4] 四种
    def __init__(self, chip_type="square", inner_chip_size=6, inter_chip_array=[3, 3], circuit_dag: Circuit_DAG = None, just_one_map: bool = False, draw_it: bool = False):
        random.seed(time.time())
        self.chip_type = chip_type
        self.inner_chip_size = inner_chip_size
        self.inter_chip_array = inter_chip_array
        self.just_one_map = just_one_map
        self.map_plan = {}
        self.draw_it = draw_it
        # 线路图，指向一个 Circuit_DAG 对象
        # 其中  circuit_dag.root_node 指向根节点
        #      circuit_dag.qubit_num 指向qubit数量
        # 需要对该DAG和量子芯片进行初始化映射
        self.circuit_dag: Circuit_DAG = circuit_dag
        # 需要综合的量子比特数量
        self.synthesis_qubit_num = circuit_dag.qubit_num
        self.chips: list[QuantumChiplet] = [
            QuantumChiplet(inner_chip_size,
                           chip_type=chip_type, chip_id=id_idx)
            for id_idx in range(inter_chip_array[0] * inter_chip_array[1])
        ]
        self.merge_all_chips_to_graph()
        self.total_qubit_num = sum([i.qubit_num for i in self.chips])
        self.map_to_chip()
        self.estimate_critical_value_with_uniform()

    def __deepcopy__(self, memo):
        '''创建一个新的 QPU 实例，不调用 __init__（避免重复初始化）'''
        new_qpu = type(self).__new__(self.__class__)
        for key, value in self.__dict__.items():
            if key in ['merged_graph', 'map_plan', 'chips', 'circuit_dag']:
                continue
            setattr(new_qpu, key, deepcopy(value, memo))
        if 'merged_graph' in self.__dict__:
            new_qpu.merged_graph = self.merged_graph.copy()
        if 'map_plan' in self.__dict__:
            new_qpu.map_plan = deepcopy(self.map_plan, memo)
        if 'chips' in self.__dict__:
            new_qpu.chips = [deepcopy(chip, memo) for chip in self.chips]
        if 'circuit_dag' in self.__dict__:
            new_qpu.circuit_dag = deepcopy(self.circuit_dag, memo)
        return new_qpu

    def move_computing_qubit(self, source_physics_qubit_id: int, target_physics_qubit_id: int):
        '''移动物理量子比特 source_physics_qubit_id 到目标的物理比特 target_physics_qubit_id 上'''
        path = nx.shortest_path(
            self.merged_graph, source_physics_qubit_id, target_physics_qubit_id)
        for i in range(len(path)-1):
            # 连续交换 path[i] 和 path[i+1]
            tmp = self.map_plan[path[i]]
            self.map_plan[path[i]] = self.map_plan[path[i+1]]
            self.map_plan[path[i+1]] = tmp

    def estimate_critical_value_with_uniform(self):
        # 不引入权重优化
        dag: nx.DiGraph = self.circuit_dag.dag_nx
        for node in dag.nodes:
            dag.nodes[node]["time_weight"] = 1

    def estimate_critical_value(self):
        # ⭐️ 根据DAG估计关键路径长度
        dag: nx.DiGraph = self.circuit_dag.dag_nx
        for node in dag.nodes:
            dag.nodes[node]["time_weight"] = 0

        def estimate_cnot_pair_distance(cnot_pair: list):
            paths = self.get_k_path_with_qubit_id(
                cnot_pair[0], cnot_pair[1], k=1)
            path = paths[0]
            return len(path)

        def get_node_weight(dag: nx.DiGraph):
            topo_reversed_order = list(
                reversed(list(nx.topological_sort(dag))))
            for node in topo_reversed_order:
                for pred in dag.predecessors(node):
                    dag.nodes[pred]["time_weight"] = max(dag.nodes[pred]["time_weight"], dag.edges[(
                        pred, node)]["weight"], dag.nodes[node]["time_weight"]+(np.log2(estimate_cnot_pair_distance(dag.nodes[node]["label"])))**2)

        dag: nx.DiGraph = self.circuit_dag.dag_nx
        get_node_weight(dag)

    def get_k_path_with_qubit_id_without_block(self, source_qubit, target_qubit, k: int = 5):
        # 在 self.map_plan 中查找值为 source 的键
        source_key = 0
        target_key = 0
        for _key, value in self.map_plan.items():
            if value == source_qubit:
                source_key = _key
            if value == target_qubit:
                target_key = _key
        return self.get_k_path_without_block(source_key, target_key, k)

    def get_k_path_with_qubit_id(self, source_qubit, target_qubit, k: int = 5):
        # 在 self.map_plan 中查找值为 source 的键
        source_key = 0
        target_key = 0
        for _key, value in self.map_plan.items():
            if value == source_qubit:
                source_key = _key
            if value == target_qubit:
                target_key = _key
        return self.get_k_path(source_key, target_key, k)

    def get_k_path_without_block(self, source, target, K: int = 5):
        # merged_graph 计算候选路径
        paths = []
        tmpG = self.merged_graph
        # 获取所有 mapped_qubits 的集合
        mapped_qubits = set()
        for key, value in self.map_plan.items():
            if value is not None:
                mapped_qubits.add(key)
        # 自定义权重函数：如果边连接了 mapped_qubit，则增加权重

        def custom_weight(u, v, d):
            weight = 1  # 默认权重
            if u in mapped_qubits or v in mapped_qubits:
                weight += 1  # 增加权重，让路径尽量避开 mapped_qubit
            return weight

        for i, path in enumerate(nx.shortest_simple_paths(tmpG, source, target, weight=custom_weight)):
            paths.append(path)
            if i == K - 1:  # 当找到 K 条路径时停止
                break
        return paths

    def get_k_path(self, source, target, K: int = 5):
        # merged_graph 计算候选路径
        paths = []
        tmpG = self.merged_graph
        # 获取所有 mapped_qubits 的集合
        mapped_qubits = set()
        for key, value in self.map_plan.items():
            if value is not None:
                mapped_qubits.add(key)
        # 自定义权重函数：如果边连接了 mapped_qubit，则增加权重

        def custom_weight(u, v, d):
            weight = 1  # 默认权重
            if u in mapped_qubits or v in mapped_qubits:
                weight += 100  # 增加权重，让路径尽量避开 mapped_qubit
            return weight
        for i, path in enumerate(nx.shortest_simple_paths(tmpG, source, target, weight=custom_weight)):
            paths.append(path)
            if i == K - 1:  # 当找到 K 条路径时停止
                break
        return paths

    def print_map_plan(self, map_plan):
        non_empty_entries = {k: v for k,
                             v in map_plan.items() if v is not None}
        print("Map Plan:", non_empty_entries)

    def map_to_chip(self):
        '''将量子比特映射到芯片'''
        def average_allocate():
            map_plan = {}
            for i in range(self.total_qubit_num):
                map_plan[i] = None
            chip_ranges = []
            current_idx = 0
            for chip in self.chips:
                chip_size = chip.qubit_num
                chip_ranges.append((current_idx, current_idx + chip_size))
                current_idx += chip_size
            synthesis_indices = list(range(self.synthesis_qubit_num))
            random.shuffle(synthesis_indices)
            total_chips = len(self.chips)
            per_chip_allocation = self.synthesis_qubit_num // total_chips
            remainder = self.synthesis_qubit_num % total_chips
            current_synthesis_idx = 0
            for i, (start, end) in enumerate(chip_ranges):
                allocation_count = per_chip_allocation + \
                    (1 if i < remainder else 0)
                chip_qubits = list(range(start, end))
                selected_physical_qubits = random.sample(
                    chip_qubits, allocation_count)
                for physical_idx in selected_physical_qubits:
                    if current_synthesis_idx >= self.synthesis_qubit_num:
                        break
                    map_plan[physical_idx] = synthesis_indices[current_synthesis_idx]
                    current_synthesis_idx += 1
            return map_plan

        def calculate_cross_cost(map_plan):
            # 计算跨芯片的代价
            # { physical_qubit: synthesis_qubit, ... }
            # 所有的CNOT门
            cx_gates = self.circuit_dag.cx_gates
            chip_ranges = []
            current_idx = 0
            total_price = 0
            for chip in self.chips:
                chip_size = chip.qubit_num
                chip_ranges.append((current_idx, current_idx + chip_size))
                current_idx += chip_size
            # for i, (start, end) in enumerate(chip_ranges):
            for cx_gate in cx_gates:
                a, b = cx_gate[0], cx_gate[1]
                physic_a, physic_b = None, None
                chip_idx_a, chip_idx_b = 0, 0
                for _key, _value in map_plan.items():
                    if (_value == a):
                        physic_a = _key
                    if (_value == b):
                        physic_b = _key
                for i, (start, end) in enumerate(chip_ranges):
                    if (physic_a >= start and physic_a < end):
                        chip_idx_a = i
                    if (physic_b >= start and physic_b < end):
                        chip_idx_b = i
                x_length = self.inter_chip_array[1]
                cross_distance = abs(int(chip_idx_a / x_length) - int(chip_idx_b / x_length)) + \
                    abs(chip_idx_a % x_length - chip_idx_b % x_length)
                total_price += cross_distance
            return total_price
        final_plan = {}
        final_price = 10 ** 10
        for _ in range(1 if self.just_one_map else 100):
            plan = average_allocate()
            price = calculate_cross_cost(plan)
            if price < final_price:
                final_plan = plan
                final_price = price
                # self.print_map_plan(plan)
        self.map_plan = final_plan
        # print("final_price:", final_price)
        if (self.draw_it):
            self.draw_chiplet()

    def merge_all_chips_to_graph(self):
        self.merged_graph = nx.Graph()
        self.global_pos = {}
        for idx, chip in enumerate(self.chips):
            for node in chip.graph.nodes():
                unique_id = chip.graph.nodes[node]['id']
                is_hexagon = "hexagon" in self.chip_type
                offset_x = idx % self.inter_chip_array[1] * (
                    self.inner_chip_size + 1) * (2.2 if is_hexagon else 1)
                offset_y = idx / self.inter_chip_array[1] * (self.inner_chip_size + 1) * (
                    1.1 if is_hexagon else 1) - (idx % self.inter_chip_array[1]) * 1.5
                self.global_pos[unique_id] = (
                    node[1] + offset_x, node[0] + offset_y)
                self.merged_graph.add_node(unique_id, id=unique_id)
            for u, v in chip.graph.edges():
                u_id = chip.graph.nodes[u]['id']
                v_id = chip.graph.nodes[v]['id']
                self.merged_graph.add_edge(u_id, v_id)
        m, n = self.inter_chip_array
        for i in range(m):
            for j in range(n):
                current_chip_id = i * n + j
                current_chip = self.chips[current_chip_id]
                if j < n - 1:
                    right_chip_id = i * n + (j + 1)
                    right_chip = self.chips[right_chip_id]
                    # 当前芯片的右边界连接右边芯片的左边界
                    indices_right = np.linspace(
                        1, len(current_chip.right_border) - 2, 3, dtype=int)
                    for idx in indices_right:
                        current_node = current_chip.right_border[idx]
                        right_node = right_chip.left_border[idx]
                        self.merged_graph.add_edge(current_node, right_node)
                if i < m - 1:
                    bottom_chip_id = (i + 1) * n + j
                    bottom_chip = self.chips[bottom_chip_id]
                    # 当前芯片的下边界连接下边芯片的上边界
                    indices_bottom = np.linspace(
                        1, len(current_chip.bottom_border) - 2, 3, dtype=int)
                    for idx in indices_bottom:
                        current_node = current_chip.bottom_border[idx]
                        bottom_node = bottom_chip.top_border[idx]
                        self.merged_graph.add_edge(current_node, bottom_node)

    def draw_chiplet(self, paths=[]):
        G = nx.Graph()
        for u, v in self.merged_graph.edges():
            G.add_edge(u, v)
        node_colors = []

        def is_start_end(node):
            flag = False
            for path in paths:
                if (node == path[0] or node == path[-1]):
                    flag = True
            return flag

        for node in G.nodes():
            if is_start_end(node):
                node_colors.append('blue')
            elif self.map_plan[node] is not None:
                node_colors.append('red')  # 被映射的节点
            else:
                node_colors.append('gray')  # 其他节点

        # plt.figure(figsize=(12, 10))  # 修改为所需的尺寸

        # 获取当前坐标轴对象并清空
        ax = plt.gca()
        ax.clear()

        pos = self.global_pos  # 假设 global_pos 是已经定义好的布局
        # 绘制整体图
        nx.draw(G, pos, with_labels=False,
                node_size=50, node_color=node_colors, ax=ax)

        # 绘制多条路径
        if paths:
            for path in paths:
                path_edges = [(path[i], path[i + 1])
                              for i in range(len(path) - 1)]
                nx.draw_networkx_edges(
                    G, pos,
                    edgelist=path_edges,
                    edge_color='red',
                    width=2,
                    arrows=True,
                    arrowstyle='-|>',  # 箭头样式
                    arrowsize=15,
                    ax=ax
                )

        plt.title("Merged Graph of All Chiplets (Mapped Qubits Highlighted)")
        plt.draw()  # 刷新当前图像
        plt.pause(0.1)  # 暂停一小段时间以保证图像更新


class QuantumChiplet:
    def __init__(self, size=4, chip_type="square", chip_id=0):
        self.chip_id = chip_id
        self.graph = None
        self.initial_idx = 0
        self.qubit_num = 0

        self.left_border = []
        self.top_border = []
        self.right_border = []
        self.bottom_border = []

        if chip_type == "square":
            self.initial_idx = chip_id * size * size
            self.qubit_num = size * size
            self.build_square_chip(size)
            self.get_square_chip_border(size)
        elif chip_type == "hexagon":
            self.initial_idx = chip_id * \
                ((2 * size + 1) * 2 + (2 * size + 2) * (size - 1))
            self.qubit_num = (2 * size + 1) * 2 + (2 * size + 2) * (size - 1)
            self.build_hexagon_chip(size)
            self.get_hexagon_chip_border(size)
        elif chip_type == "heavy_square":
            self.initial_idx = chip_id * (size * (3 * size - 2))
            self.qubit_num = size * (3 * size - 2)
            self.build_heavy_square_chip(size)
            self.get_heavy_square_chip_border(size)
        elif chip_type == "heavy_hexagon":
            self.initial_idx = chip_id * \
                ((4 * size + 1) * 2 + size * (size + 1) + (4 * size + 3) * (size - 1))
            self.qubit_num = (4 * size + 1) * 2 + size * \
                (size + 1) + (4 * size + 3) * (size - 1)
            self.build_heavy_hexagon_chip(size)
            self.get_heavy_hexagon_chip_border(size)
        else:
            raise ValueError(f"Unsupported chip type: {chip_type}")
        # self.visualize()

    def get_square_chip_border(self, size):
        self.left_border = [
            self.graph.nodes[(i, 0)]['id'] for i in range(size)]
        self.right_border = [
            self.graph.nodes[(i, size-1)]['id'] for i in range(size)]

        self.top_border = [self.graph.nodes[(0, i)]['id'] for i in range(size)]
        self.bottom_border = [
            self.graph.nodes[(size-1, i)]['id'] for i in range(size)]

    def get_hexagon_chip_border(self, size):
        # return
        is_odd = size % 2 == 1
        self.left_border = [self.graph.nodes[(i, 0)]['id'] for i in range(
            size)] + [self.graph.nodes[(size, 1)]['id']]
        self.right_border = [self.graph.nodes[(0, 2*size)]['id']] + [
            self.graph.nodes[(i+1, 2*size+1)]['id'] for i in range(size if not is_odd else size-1)] + ([self.graph.nodes[(size, 2*size)]['id']] if is_odd else [])
        self.top_border = [
            self.graph.nodes[(0, 2*i+1)]['id'] for i in range(size)]
        self.bottom_border = [
            self.graph.nodes[(size, 2*i+2)]['id'] for i in range(size)]

    def get_heavy_square_chip_border(self, size):
        self.get_square_chip_border(size)

    def get_heavy_hexagon_chip_border(self, size):
        self.get_hexagon_chip_border(size)

    def build_square_chip(self, size):
        self.graph = nx.grid_2d_graph(size, size)
        for local_idx, node in enumerate(self.graph.nodes()):
            unique_id = self.initial_idx + local_idx
            self.graph.nodes[node]['id'] = unique_id

    def build_hexagon_chip(self, size):
        width, height = size, size
        self.graph = nx.hexagonal_lattice_graph(width, height)
        for local_idx, node in enumerate(self.graph.nodes()):
            unique_id = self.initial_idx + local_idx
            self.graph.nodes[node]['id'] = unique_id

    def build_heavy_square_chip(self, size):
        G = nx.grid_2d_graph(size, size)
        mid_edge_nodes = []
        edges_list = list(G.edges())
        for (u, v) in edges_list:
            x1, y1 = u
            x2, y2 = v
            mid_node = ((x1 + x2)/2, (y1 + y2)/2)
            mid_edge_nodes.append(mid_node)
            G.add_node(mid_node)
            G.add_edge(u, mid_node)
            G.add_edge(v, mid_node)

        self.graph = G
        for local_idx, node in enumerate(self.graph.nodes()):
            unique_id = self.initial_idx + local_idx
            self.graph.nodes[node]['id'] = unique_id

    def build_heavy_hexagon_chip(self, size):
        width, height = size, size
        G = nx.hexagonal_lattice_graph(width, height)
        edges_list = list(G.edges())
        for (u, v) in edges_list:
            x1, y1 = u
            x2, y2 = v
            mid_x = (x1 + x2) / 2
            mid_y = (y1 + y2) / 2
            mid_node = (mid_x, mid_y)
            G.remove_edge(u, v)
            G.add_node(mid_node)
            G.add_edge(u, mid_node)
            G.add_edge(mid_node, v)
        self.graph = G
        for local_idx, node in enumerate(self.graph.nodes()):
            unique_id = self.initial_idx + local_idx
            self.graph.nodes[node]['id'] = unique_id

    def visualize(self):
        plt.figure(figsize=(6, 6))
        pos = {(x, y): (y, -x) for (x, y) in self.graph.nodes()}
        labels = {}
        for node in self.graph.nodes():
            x, y = node
            if x == int(x) and y == int(y):
                labels[node] = str(node)
        nx.draw(self.graph, pos=pos, with_labels=False,
                node_size=50, font_size=8)
        nx.draw_networkx_labels(self.graph, pos=pos,
                                labels=labels, font_size=8)
        plt.title(f"{type(self).__name__} Graph")
        plt.show()
